import struct
import numpy as np

# Error codes matching the C version
WAV_OK = 0
WAV_ERROR = -1
WAV_NO_MEMORY = -2
WAV_INVALID_ARG = -3
WAV_IO_ERROR = -4

# ADPCM tables matching the C version
STEP_TABLE = [
    7, 8, 9, 10, 11, 12, 13, 14, 16, 17, 19, 21, 23, 25, 28, 31, 34, 37, 41, 45,
    50, 55, 60, 66, 73, 80, 88, 97, 107, 118, 130, 143, 157, 173, 190, 209, 230,
    253, 279, 307, 337, 371, 408, 449, 494, 544, 598, 658, 724, 796, 876, 963,
    1060, 1166, 1282, 1411, 1552, 1707, 1878, 2066, 2272, 2499, 2749, 3024, 3327,
    3660, 4026, 4428, 4871, 5358, 5894, 6484, 7132, 7845, 8630, 9493, 10442, 11487,
    12635, 13899, 15289, 16818, 18500, 20350, 22385, 24623, 27086, 29794, 32767
]

INDEX_TABLE = [
    -1, -1, -1, -1, 2, 4, 6, 8, -1, -1, -1, -1, 2, 4, 6, 8
]

class CodecConfig:
    """Compression configuration matching the C version"""
    def __init__(self):
        self.target_sample_rate = 0  # 0 means keep original sample rate
        self.target_bits = 8         # Reduce to 8 bits
        self.use_adpcm = 1           # Use ADPCM encoding

def codec_init_config(config):
    """Initialize default compression configuration"""
    if config is None:
        return
    
    config.target_sample_rate = 0  # 0 means keep original sample rate
    config.target_bits = 8         # Reduce to 8 bits
    config.use_adpcm = 1           # Use ADPCM encoding

def read_wav_header(buffer):
    """Read WAV header from buffer"""
    if len(buffer) < 44:  # WAV header size
        return None, WAV_INVALID_ARG
    
    header = {}
    header['riff_id'] = buffer[0:4]
    header['size'] = struct.unpack('<I', buffer[4:8])[0]
    header['wave_id'] = buffer[8:12]
    header['fmt_id'] = buffer[12:16]
    header['fmt_size'] = struct.unpack('<I', buffer[16:20])[0]
    header['format'] = struct.unpack('<H', buffer[20:22])[0]
    header['channels'] = struct.unpack('<H', buffer[22:24])[0]
    header['sample_rate'] = struct.unpack('<I', buffer[24:28])[0]
    header['byte_rate'] = struct.unpack('<I', buffer[28:32])[0]
    header['block_align'] = struct.unpack('<H', buffer[32:34])[0]
    header['bits_per_sample'] = struct.unpack('<H', buffer[34:36])[0]
    header['data_id'] = buffer[36:40]
    header['data_size'] = struct.unpack('<I', buffer[40:44])[0]
    
    # Verify WAV format
    if (header['riff_id'] != b'RIFF' or 
        header['wave_id'] != b'WAVE' or 
        header['fmt_id'] != b'fmt '):
        return None, WAV_ERROR
    
    # Only support PCM format
    if header['format'] != 1:
        return None, WAV_ERROR
    
    return header, WAV_OK

def downsample(input_data, factor):
    """Reduce sampling rate by taking every nth sample"""
    return input_data[::factor]

def reduce_bit_depth(input_data, target_bits):
    """Reduce bit depth of 16-bit samples to target_bits"""
    shift = 16 - target_bits
    max_value = (1 << target_bits) - 1
    
    output = np.zeros(len(input_data), dtype=np.uint8)
    for i in range(len(input_data)):
        # Convert signed 16-bit to unsigned target_bits
        scaled = ((input_data[i] + 32768) >> shift) & max_value
        
        # Ensure in range
        if scaled > max_value:
            scaled = max_value
        if scaled < 0:
            scaled = 0
            
        output[i] = scaled
    
    return output

def encode_adpcm_sample(sample, predicted, index):
    """Encode a single sample using ADPCM"""
    step = STEP_TABLE[index]
    diff = sample - predicted
    code = 0
    
    if diff < 0:
        code = 8
        diff = -diff
    
    tempstep = step
    if diff >= tempstep:
        code |= 4
        diff -= tempstep
    
    tempstep >>= 1
    if diff >= tempstep:
        code |= 2
        diff -= tempstep
    
    tempstep >>= 1
    if diff >= tempstep:
        code |= 1
    
    # Adjust predicted value
    diff = step >> 3
    if code & 4:
        diff += step
    if code & 2:
        diff += step >> 1
    if code & 1:
        diff += step >> 2
    
    if code & 8:
        predicted -= diff
        if predicted < -32768:
            predicted = -32768
    else:
        predicted += diff
        if predicted > 32767:
            predicted = 32767
    
    # Adjust index
    index += INDEX_TABLE[code & 15]
    if index < 0:
        index = 0
    if index > 88:
        index = 88
    
    return code, predicted, index

def decode_adpcm_sample(code, predicted, index):
    """Decode a single sample using ADPCM"""
    step = STEP_TABLE[index]
    diff = step >> 3
    
    if code & 4:
        diff += step
    if code & 2:
        diff += step >> 1
    if code & 1:
        diff += step >> 2
    
    if code & 8:
        predicted -= diff
        if predicted < -32768:
            predicted = -32768
    else:
        predicted += diff
        if predicted > 32767:
            predicted = 32767
    
    index += INDEX_TABLE[code & 15]
    if index < 0:
        index = 0
    if index > 88:
        index = 88
    
    return predicted, index

def adpcm_encode(input_data):
    """Encode audio data using ADPCM"""
    samples = len(input_data)
    out_size = (samples + 1) // 2  # Each byte stores two 4-bit samples
    output = bytearray(out_size)
    
    predicted = 0
    index = 0
    out_index = 0
    
    # Each byte can store two ADPCM samples (4 bits each)
    for i in range(0, samples, 2):
        code1, predicted, index = encode_adpcm_sample(input_data[i], predicted, index)
        
        if i + 1 < samples:
            code2, predicted, index = encode_adpcm_sample(input_data[i+1], predicted, index)
            output[out_index] = (code1 << 4) | code2
        else:
            output[out_index] = (code1 << 4)
        
        out_index += 1
    
    return bytes(output)

def adpcm_decode(input_data, output_samples):
    """Decode ADPCM audio data"""
    encoded_size = len(input_data)
    output = np.zeros(output_samples, dtype=np.int16)
    
    predicted = 0
    index = 0
    out_index = 0
    
    for i in range(encoded_size):
        if out_index >= output_samples:
            break
            
        code1 = (input_data[i] >> 4) & 0x0F
        predicted, index = decode_adpcm_sample(code1, predicted, index)
        output[out_index] = predicted
        out_index += 1
        
        if out_index >= output_samples:
            break
            
        code2 = input_data[i] & 0x0F
        predicted, index = decode_adpcm_sample(code2, predicted, index)
        output[out_index] = predicted
        out_index += 1
    
    return output

def wav_compress(in_buffer, config):
    """Compress WAV data
    
    Args:
        in_buffer (bytes): Input WAV data including header
        config (CodecConfig): Compression configuration
        
    Returns:
        tuple: (compressed_data, result_code)
    """
    if in_buffer is None or config is None:
        return None, WAV_INVALID_ARG
    
    # Parse WAV header
    header, result = read_wav_header(in_buffer)
    if result != WAV_OK:
        return None, result
    
    # Initialize compressed header (matching C struct)
    comp_header = {
        'original_size': len(in_buffer),
        'sample_rate': header['sample_rate'],
        'channels': header['channels'],
        'original_bits': header['bits_per_sample'],
        'compressed_size': 0,
        'target_sample_rate': 0,
        'target_bits': config.target_bits,
        'use_adpcm': config.use_adpcm
    }
    
    # Determine target sample rate
    if config.target_sample_rate == 0:
        comp_header['target_sample_rate'] = header['sample_rate']  # Keep original rate
    else:
        comp_header['target_sample_rate'] = config.target_sample_rate
    
    # Calculate data offset and extract audio data
    data_offset = 44  # WAV header size
    audio_data = np.frombuffer(in_buffer[data_offset:], dtype=np.int16)
    
    # Use PCM compression function to process audio data
    compressed_data, result = pcm_compress(audio_data, config)
    if result != WAV_OK:
        return None, result
    
    comp_header['compressed_size'] = len(compressed_data)
    
    # Create output buffer with header (24 bytes to match C struct)
    # Format: '<IIHHIIBC' for exact compatibility with C struct CompressedHeader
    out_buffer = bytearray()
    out_buffer.extend(struct.pack('<I', comp_header['original_size']))       # unsigned int original_size
    out_buffer.extend(struct.pack('<I', comp_header['sample_rate']))         # unsigned int sample_rate
    out_buffer.extend(struct.pack('<H', comp_header['channels']))            # unsigned short channels
    out_buffer.extend(struct.pack('<H', comp_header['original_bits']))       # unsigned short original_bits
    out_buffer.extend(struct.pack('<I', comp_header['compressed_size']))     # unsigned int compressed_size
    out_buffer.extend(struct.pack('<I', comp_header['target_sample_rate']))  # unsigned int target_sample_rate
    out_buffer.extend(struct.pack('<H', comp_header['target_bits']))         # unsigned short target_bits
    out_buffer.extend(struct.pack('<B', comp_header['use_adpcm']))          # unsigned char use_adpcm
    
    # Add padding to match C struct alignment (CompressedHeader)
    out_buffer.extend(bytes(3))  # Add 3 bytes of padding to make it 24 bytes total
    
    # Add compressed data
    out_buffer.extend(compressed_data)
    
    return bytes(out_buffer), WAV_OK

def wav_decompress(in_buffer):
    """Decompress WAV data
    
    Args:
        in_buffer (bytes): Input compressed data
        
    Returns:
        tuple: (decompressed_data, result_code)
    """
    if in_buffer is None or len(in_buffer) < 24:  # CompressedHeader size
        return None, WAV_INVALID_ARG
    
    # Read compressed header
    comp_header = {}
    comp_header['original_size'] = struct.unpack('<I', in_buffer[0:4])[0]
    comp_header['sample_rate'] = struct.unpack('<I', in_buffer[4:8])[0]
    comp_header['channels'] = struct.unpack('<H', in_buffer[8:10])[0]
    comp_header['original_bits'] = struct.unpack('<H', in_buffer[10:12])[0]
    comp_header['compressed_size'] = struct.unpack('<I', in_buffer[12:16])[0]
    comp_header['target_sample_rate'] = struct.unpack('<I', in_buffer[16:20])[0]
    comp_header['target_bits'] = struct.unpack('<H', in_buffer[20:22])[0]
    comp_header['use_adpcm'] = in_buffer[22]
    
    # Calculate expected decompressed samples
    if comp_header['use_adpcm']:
        expected_samples = comp_header['compressed_size'] * 2
    else:
        expected_samples = comp_header['compressed_size']
    
    # Use PCM decompression function to process audio data
    decompressed, result = pcm_decompress(
        in_buffer[24:24+comp_header['compressed_size']], 
        expected_samples,
        comp_header['use_adpcm'], 
        comp_header['target_bits']
    )
    
    if result != WAV_OK:
        return None, result
    
    # Create WAV header
    data_size = len(decompressed) * 2  # 2 bytes per sample (16-bit)
    wav_header = bytearray()
    wav_header.extend(b'RIFF')
    wav_header.extend(struct.pack('<I', 36 + data_size))  # File size - 8
    wav_header.extend(b'WAVE')
    wav_header.extend(b'fmt ')
    wav_header.extend(struct.pack('<I', 16))  # fmt chunk size
    wav_header.extend(struct.pack('<H', 1))   # PCM format
    wav_header.extend(struct.pack('<H', comp_header['channels']))
    wav_header.extend(struct.pack('<I', comp_header['target_sample_rate']))
    byte_rate = comp_header['target_sample_rate'] * comp_header['channels'] * 2  # 2 bytes per sample
    wav_header.extend(struct.pack('<I', byte_rate))
    block_align = comp_header['channels'] * 2  # 2 bytes per sample
    wav_header.extend(struct.pack('<H', block_align))
    wav_header.extend(struct.pack('<H', 16))  # 16-bit after decompression
    wav_header.extend(b'data')
    wav_header.extend(struct.pack('<I', data_size))
    
    # Create output buffer
    out_buffer = bytearray(wav_header)
    out_buffer.extend(decompressed.tobytes())
    
    return bytes(out_buffer), WAV_OK

def pcm_compress(in_data, config):
    """Compress PCM audio data without WAV header processing
    
    Args:
        in_data (numpy.ndarray): Input PCM data as 16-bit signed integers
        config (CodecConfig): Compression configuration
        
    Returns:
        tuple: (compressed_data, result_code)
    """
    if in_data is None or config is None:
        return None, WAV_INVALID_ARG
    
    # Determine downsample factor
    downsample_factor = 1  # Default to no downsampling
    
    if config.target_sample_rate != 0 and config.target_sample_rate < 44100:
        # Only calculate downsample factor if we need to reduce sample rate
        downsample_factor = max(1, 44100 // config.target_sample_rate)
    
    # Downsample audio or keep original
    if downsample_factor > 1:
        downsampled = downsample(in_data, downsample_factor)
    else:
        downsampled = in_data  # Keep original samples
    
    if config.use_adpcm:
        # ADPCM encode
        compressed_data = adpcm_encode(downsampled)
        return compressed_data, WAV_OK
    else:
        # Reduce bit depth
        reduced_bits = reduce_bit_depth(downsampled, config.target_bits)
        return bytes(reduced_bits), WAV_OK

def pcm_decompress(in_buffer, out_samples, use_adpcm, target_bits):
    """Decompress PCM audio data without WAV header processing
    
    Args:
        in_buffer (bytes): Input compressed data
        out_samples (int): Expected number of output samples
        use_adpcm (bool): Whether ADPCM encoding was used
        target_bits (int): Target bit depth used for compression
        
    Returns:
        tuple: (decompressed_data, result_code)
    """
    if in_buffer is None or out_samples <= 0:
        return None, WAV_INVALID_ARG
    
    if use_adpcm:
        # ADPCM decode
        decompressed = adpcm_decode(in_buffer, out_samples)
    else:
        # Convert reduced bit depth data back to 16-bit
        decompressed = np.zeros(out_samples, dtype=np.int16)
        
        for i in range(min(out_samples, len(in_buffer))):
            # Convert from unsigned to signed properly
            shift = 16 - target_bits
            
            # Scale from unsigned target_bits back to signed 16-bit
            unsigned_val = in_buffer[i]
            signed_val = (unsigned_val << shift) - 32768
            decompressed[i] = signed_val
    
    return decompressed, WAV_OK 